import os

from neo4j import GraphDatabase

from embedding import get_embedding

GRAPH_HOST=os.getenv("GRAPH_HOST")
GRAPH_PORT=os.getenv("GRAPH_PORT")
GRAPH_USER=os.getenv("GRAPH_USER")
GRAPH_PASSWORD=os.getenv("GRAPH_PASSWORD")
GRAPH_DATABASE=os.getenv("GRAPH_DATABASE")

def main():
    driver = GraphDatabase.driver(f"bolt://{GRAPH_HOST}:{GRAPH_PORT}", auth=(GRAPH_USER, GRAPH_PASSWORD), database=GRAPH_DATABASE)
    # 定义要处理的节点的标签
    node_labels = ["Book", "Author"]
    for node_label in node_labels:
        result = driver.execute_query(f"MATCH (n:{node_label}) RETURN n.name as name")
        nodes = result.records
        print(f"共找到{len(nodes)}个{node_label}节点")
        for record in nodes:
            data = record.data()
            name = data.get("name")
            embedding = get_embedding(name)
            driver.execute_query(
                f"MATCH (n:{node_label}) where n.name = $name set n.embedding=$embedding",
                name=name,
                embedding=embedding
            )
    print("标签向量化完成")
    driver.close()

if __name__ == "__main__":
    main()